import math
from math import ceil, sqrt, floor
from random import random

import torch

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from individual.population import Population
from problems.problem import Problem
from utils.sort import sortrows
from utils.vectors import get_ref_vectors, vector_radian


class FLEA(Algorithm):
    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.population = problem.init_population(pop_size)
        self.RN = ceil(sqrt(pop_size))

    def each_iteration(self):
        front_no, _ = NSGA2.ND_sort(self.population.obj, len(self.population))
        Ref = self.ref_selection(self.problem, self.population, self.RN,
                                 (self.problem.fun_eval / self.problem.max_fun_eval) ** 2)
        [id_obj, id_dec] = self.neighborhood_association(self.population, Ref)
        cm = 0.1 * floor(10 * (self.problem.fun_eval / self.problem.max_fun_eval))
        dirV_obj = self.direction_calculation(self.problem, self.population, Ref, id_obj, front_no)
        dirV_dec = self.direction_calculation(self.problem, self.population, Ref, id_dec, front_no)
        offspring = self.reproduction(self.problem, Ref, cm, dirV_obj, dirV_dec)
        temp_pop = self.population + offspring
        self.population = temp_pop[NSGA2.environmental_selection(temp_pop.obj, self.pop_size)]

    @staticmethod
    def ref_selection(problem: Problem, population, RN, theta):
        v = get_ref_vectors(RN, problem.obj_dim)
        pop_obj = population.obj
        v = v * (pop_obj.max(0)[0] - pop_obj.min(0)[0] + math.exp(1))
        N, M = pop_obj.shape
        NV = len(v)
        pop_obj = pop_obj - pop_obj.min(0)[0]
        gamma = vector_radian(v, v, return_radian=True).fill_diagonal_(torch.inf).min(1)[0]
        Angle = vector_radian(pop_obj, v, return_radian=True)
        next_index = torch.zeros(NV, dtype=torch.long)
        for i in range(NV):
            APD = (1 + problem.obj_dim * theta * Angle[:, i] / gamma[i]) * torch.sqrt(torch.sum(pop_obj ** 2, 1))
            next_index[i] = torch.argmin(APD)
            Angle[next_index[i]] = torch.inf
        Ref = population[next_index]
        _, order1 = sortrows(Ref.obj, return_sorted_indices=True)
        return Ref[order1]

    @staticmethod
    def neighborhood_association(population, Ref):
        RefObj = Ref.obj
        PopObj = population.obj
        Distance_Obj = torch.max(torch.abs(RefObj.unsqueeze(1) - PopObj.unsqueeze(0)), dim=2)[0]
        OCid = torch.min(Distance_Obj, dim=0)[1]
        RefDec = Ref.dec
        PopDec = population.dec
        Distance_Dec = torch.max(torch.abs(RefDec.unsqueeze(1) - PopDec.unsqueeze(0)), dim=2)[0]
        DCid = torch.min(Distance_Dec, dim=0)[1]
        return OCid, DCid

    @staticmethod
    def direction_calculation(problem, population, Ref, id1, front_no):
        RefDec = Ref.dec
        RN = len(Ref)
        dirV = torch.zeros_like(RefDec)
        for i in range(RN):
            Neighbor = population[id1 == i]
            F_Neighbor = front_no[id1 == i]
            if len(Neighbor) > 0 and min(F_Neighbor) > 1:
                dirV[i] = torch.mean(Neighbor[F_Neighbor == min(F_Neighbor)].dec, dim=0) - RefDec[i]
            elif random() > 0.5:
                dirV[i] = (RefDec[i] - problem.low_limit) / problem.var_dim
            else:
                dirV[i] = (problem.high_limit - RefDec[i]) / problem.var_dim
        return dirV

    @staticmethod
    def reproduction(problem, Ref, cm, dirV_obj, dirV_dec):
        RefDec = Ref.dec
        RN = len(Ref)
        CN = floor(cm * 0.5 * problem.pop_size / RN)
        DN = problem.pop_size - CN * RN * 2
        OffDec = torch.zeros(CN * RN * 2 + DN, problem.var_dim)
        for i in range(RN):
            mu = RefDec[i]
            OffDec[i * CN: (i + 1) * CN] = mu + torch.randn((CN, 1), dtype=torch.double) * dirV_obj[i]
            OffDec[CN * RN + i * CN: CN * RN + (i + 1) * CN] = mu + torch.randn((CN, 1), dtype=torch.double) * \
                                                               dirV_dec[i]
        x_1 = torch.randint(0, RN, (DN,))
        x_2 = torch.randint(0, RN, (DN,))
        try:
            dirV = Ref[x_1].dec - Ref[x_2].dec
            OffDec[problem.pop_size - DN:] = Ref[x_1].dec + torch.randn((DN, 1), dtype=torch.double) * dirV
        except RuntimeError:
            return Population()
        return problem.estimate_population(OffDec)
